#!/usr/bin/env python

from itertools import product
import gurobipy as grb

def parse_lines(iterator, rows, cols):
    """parse cols*rows lines from iterator.
    return list of lists where
    xs[row] = [ids of cols with value 1]
    ex.:
    input
    0
    0
    1
    1
    parse_lines(_, 2, 2):
    xs = [[], [0, 1]]
    """
    xs = [[] for _ in xrange(rows)]
    for row in xrange(rows):
        for col in xrange(cols):
            val = int(next(iterator))
            if val == 1:
                xs[row].append(col)
    return xs


def add_indicator(model, variable, value, bigM, aux=None, indicator=None):
    import gurobipy as grb
    update = False
    if aux is None:
        aux = model.addVar(vtype=grb.GRB.BINARY)
        update = True
    if indicator is None:
        indicator = model.addVar(vtype=grb.GRB.BINARY)
        update = True

    if update:
        model.update()

    model.addConstr(variable <= value - 1 + bigM * aux + indicator)
    model.addConstr(variable >= value + 1 - (1 - aux) * bigM)
    return indicator


def mycallback(model, where):
    if where == grb.GRB.callback.MIPSOL:
        obj = model.cbGet(grb.GRB.callback.MIPSOL_OBJ)
        nodecnt = int(model.cbGet(grb.GRB.callback.MIPSOL_NODCNT))
        print '*** New solution at node', nodecnt, 'objective', obj
        #print model.cbGetSolution(model.getVars())
        with open(model._filename + str(m._solcounter) + ".sol", "w") as f:
            f.write("obj = {0}\n".format(obj))
            for (index, var) in model._variables.items():
                x = model.cbGetSolution(var)
                if x > 0.5:
                    f.write("{0}\n".format(index, var))
        m._solcounter += 1




class Timetable:
    def __init__(self, room_sizes, student_event, room_feature, event_feature,
                    num_features):
        # save some information
        self.rooms = range(len(room_sizes))
        self.room_sizes = room_sizes
        self.students = range(len(student_event))
        self.events = range(len(event_feature))
        self.features = range(num_features)

        # create mappings for student <-> event
        event_student = [[] for _ in self.events]
        for student, events in enumerate(student_event):
            for event in events:
                event_student[event].append(student)
        self.student_event = student_event
        self.event_student = event_student

        # create mappings for room <-> event pairings
        room_event = [[] for _ in self.rooms]
        event_room = [[] for _ in self.events]
        for room, event in product(self.rooms, self.events):
            roomsize = room_sizes[room]
            num_participants = len(event_student[event])
            if roomsize < num_participants:
                continue
            if not set(event_feature[event]).issubset(
                set(room_feature[room])):
                continue
            room_event[room].append(event)
            event_room[event].append(room)
        self.event_room = event_room
        self.room_event = room_event

        # save the rest. these values are not needed but could be used for
        # feasibility checking
        self.room_feature = room_feature
        self.event_feature = event_feature

        # default settings:
        # 5 days a 9 slots
        self.days = range(5)
        self.hours = range(9)

        self.timeslots = [(d, h) for d in self.days for h in self.hours]

    def __repr__(self):
        return """<Timetable {0} students, {1} rooms, """.format(
                len(self.students), len(self.rooms)) + \
               """{0} events, {1} features>""".format(
                len(self.events), len(self.features))

    def generate_ip(self, verbose=True, 
                        single_event_penalty=True,
                        last_hour_penalty=True):
        import gurobipy as grb

        if verbose:
            print "generating model"
        m = grb.Model(self.__repr__())


        if verbose:
            print "adding variables"
        # lets go
        variables = dict()
        for h in self.hours:
            for e in self.events:
                obj = 0
                if last_hour_penalty and h == self.hours[-1]:
                    obj = len(self.event_student[e])
                for d in self.days:
                    t = (d, h)
                    for r in self.rooms:
                        variables[(e, r, t)] = m.addVar(vtype=grb.GRB.BINARY,
                        obj=obj, name="e{0}_r{1}_d{2}_h{3}".format(e, r, *t))

        # adding #day x #students many auxiliary and indicator variables
        aux_variables = dict(((student, day),
                                m.addVar(vtype=grb.GRB.BINARY,
                                name="aux_s{0}_d{1}".format(student, day)))
                            for student, day in product(self.students, self.days))
        indicator_variables = dict(((student, day),
                                m.addVar(vtype=grb.GRB.BINARY,
                                name="ind_s{0}_d{1}".format(student, day)))
                            for student, day in product(self.students, self.days))

        #for e in self.events:
        #    num_students = len(self.event_student[e])
        #    for r in self.rooms:
        #        for t
        #variables = dict(((e, r, t), m.addVar(vtype=grb.GRB.BINARY,
        #    obj=()
        #    name="e{0}_r{1}_d{2}_h{3}".format(e, r, *t)))
        #    for (e, r, t) in all_ind())
        m.update()


        if verbose:
            print "adding event<->room pairing constraints"
        # restrict room <-> event pairing
        for (e, r) in product(self.events, self.rooms):
            for t in self.timeslots:
                if e in self.room_event[r]:
                    continue
                m.addConstr(variables[(e, r, t)] == 0)
        #m.update()


        if verbose:
            print "adding non-overlapping events constraints"
        # at most one event per room & timeslot
        for (r, t) in product(self.rooms, self.timeslots):
            m.addConstr(grb.quicksum(variables[(e, r, t)]
                for e in self.room_event[r]) <= 1)
        #m.update()


        if verbose:
            print "adding student/event constraints"
        # every student goes to at most one event at the same time
        for student in self.students:
            for t in self.timeslots:
                m.addConstr(grb.quicksum(variables[(e, r, t)]
                    for e in self.student_event[student]
                    for r in self.event_room[e]) <= 1)
        #m.update()


        if verbose:
            print "adding event<->timeslot/room constraints"
        # find slot / room for every event
        for e in self.events:
            m.addConstr(grb.quicksum(variables[(e, r, t)]
                for r in self.event_room[e]
                for t in self.timeslots) == 1)
        #m.update()

        # add objective functions
        # a student has a class in the last slot of the day;
        #for day in self.days:
        #    t = (day, self.hours[-1])
        #    for (e, r) in product(self.events, self.rooms):
        #        variables[(e, r, t)].obj = len(self.event_student[e])


        if single_event_penalty:
            if verbose:
                print "adding single event penalty"
            # a student has a single class on a day.
            # sum up the number of classes on each day for each student
            # if the number is == 1: incur fixed cost
            num_classes = dict()
            for s in self.students:
                if len(self.hours) == 1:
                    print "only one slot per day, check penalties!"
                for day in self.days:
                    num_classes = \
                        grb.quicksum(variables[(e, r, (day, t))]
                        for e in self.student_event[s]
                        for r in self.rooms
                        for t in self.hours)
                    indicator = add_indicator(m, num_classes, 1, len(self.hours) + 1,
                        aux=aux_variables[(s, day)],
                        indicator=indicator_variables[(s, day)])
                    indicator.obj = 1
        print "updating ...",
        m.update()

        # adding information for callback
        m._variables = variables
        #m._instance = self
        print "done"
        m.params.heuristics = 0.15
        m._solcounter = 0
        # a student has more than two classes consecutively;
        return m

    @staticmethod
    def from_file(filename):
        with open(filename) as f:
            line_iter = iter(f)

            # parse meta information
            meta = line_iter.next()
            num_events, num_rooms, num_features, num_students = \
                map(int, meta.split())

            # parse room sizes
            room_sizes = [int(line_iter.next()) for _ in xrange(num_rooms)]

            # parse student/event matrix
            student_event = parse_lines(line_iter, num_students, num_events)

            # parse room/feature matrix
            room_feature = parse_lines(line_iter, num_rooms, num_features)

            # parse event/feature matrix
            event_feature = parse_lines(line_iter, num_events, num_features)

            iteratorempty = False
            try:
                line_iter.next()
            except StopIteration:
                iteratorempty = True
            if iteratorempty is False:
                print "possible bug, not all information consumed!"
                import IPython
                IPython.embed()

            assert(num_students == len(student_event))
            assert(num_rooms == len(room_sizes))
            assert(num_events == len(event_feature))

            #import IPython
            #IPython.embed()
        return Timetable(room_sizes=room_sizes,
                         student_event=student_event,
                         room_feature=room_feature,
                         event_feature=event_feature,
                         num_features=num_features)


def timetableEA(instance):
    from pyevolve import G1DBinaryString, GSimpleGA
    #eval_func = lambda s: instance.num_satisfied_clauses(s)

    def eval_func(s, instance=instance):
        return instance.num_satisfied_clauses(s.genomeList)

    #eval_func = partial(eval_func, instance=instance)

    # Genome instance
    genome = G1DBinaryString.G1DBinaryString(instance.num_variables)

    # The evaluator function (objective function)
    genome.evaluator.set(eval_func)
    ga = GSimpleGA.GSimpleGA(genome)

    # some paramters
    ga.setGenerations(1000)
    #ga.setCrossoverRate(1.0)
    ga.setMutationRate(0.1)
    ga.setPopulationSize(100)

    # Do the evolution, with stats dump
    # frequency of 10 generations
    ga.setInteractiveGeneration(1)
    ga.evolve(freq_stats=10)

    # Best individual
    print ga.bestIndividual()

if __name__ == '__main__':
    import sys
    import IPython
    filename = "data/Timetabling/simple_test.tim"
    if len(sys.argv) > 1:
        filename = sys.argv[1]
    simple = Timetable.from_file(filename)
    m = simple.generate_ip()
    m._filename = filename
    m.optimize(mycallback)
    m.printAttr("X")
    IPython.embed()
